
import os, glob
import torch
from torch import nn
from torch import optim
from torch.utils.data import DataLoader
from PIL import Image
from torchvision import transforms
from torch import Tensor
import csv, time
import pickle

def getFeatureMap(model, device, imgPath:str, savePath:str=None, hasRet=False):

    if savePath is not None:
        savePath = savePath + "/" + imgPath.split("/")[-1][:-3] + "pt"
        if os.path.exists(savePath):
            return
        #     if hasRet:
        #         return torch.load(hasRet)
        #     else:
        #         return

    tf = transforms.Compose([
        lambda x: Image.open(x).convert('RGB'),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        # transforms.Normalize(mean=[0.485, 0.456, 0.406],
        #                     std=[0.229, 0.224, 0.225])
    ])

    x = tf(imgPath)
    x = x.unsqueeze(0).to(torch.device(device))
    # y = globals()[modelName](x)
    with torch.no_grad():
        y = model(x)
    y = y.squeeze(0)
    # print(y.shape)

    # print(img)
    # print(".\\data\\VeRi\\image_test_featuremap\\" + img.split("\\")[-1][:-3] + "pt")
    if savePath is not None:
        torch.save(y, savePath)
        # print("save: ", savePath)

    if hasRet:
        return y
    else:
        return

def getFeatureMapFromArray(model, device, imgArray):

    tf = transforms.Compose([
        lambda x: Image.fromarray(x).convert('RGB'),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        # transforms.Normalize(mean=[0.485, 0.456, 0.406],
        #                     std=[0.229, 0.224, 0.225])
    ])

    x = tf(imgArray)
    x = x.unsqueeze(0).to(torch.device(device))
    # y = globals()[modelName](x)
    with torch.no_grad():
        y = model(x)
    y = y.squeeze(0)
    return y

# 获取列表图片的特征图
def getFeatureMaps(model, device, path):

    tf = transforms.Compose([
        lambda x: Image.open(x).convert('RGB'),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        # transforms.Normalize(mean=[0.485, 0.456, 0.406],
        #                     std=[0.229, 0.224, 0.225])
    ])

    imgsName = glob.glob(path + "/*.jpg")
    # print(imgsName)
    imgs = [tf(i).unsqueeze(0) for i in imgsName]
    imgs = torch.cat(imgs, dim=0).to(torch.device(device))
    # print(imgs.shape)
    with torch.no_grad():
        y = model(imgs)
    for i in range(len(imgsName)):
        # savePath = path + "/" + imgsName[i].split("/")[-1][:-3] + "pt"
        savePath = "{0}/{1}pt".format(path, os.path.basename(imgsName[i])[:-3])
        torch.save(y[i], savePath)


def getFeatureMapsToFMList(model, device, path):

    tf = transforms.Compose([
        lambda x: Image.open(x).convert('RGB'),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        # transforms.Normalize(mean=[0.485, 0.456, 0.406],
        #                     std=[0.229, 0.224, 0.225])
    ])

    imgsName = glob.glob(path + "/*.jpg")
    # print(imgsName)
    imgs = [tf(i).unsqueeze(0) for i in imgsName]
    imgs = torch.cat(imgs, dim=0).to(torch.device(device))
    # print(imgs.shape)
    with torch.no_grad():
        y = model(imgs)

    return [[int(os.path.basename(imgsName[i])[3:-4]), y[i]] for i in range(len(imgsName))]


def getFeatureMapsFromDumpsToFMList(model, device, data):

    tf = transforms.Compose([
        lambda x: Image.fromarray(x).convert('RGB'),
        # transforms.Resize((224, 224)),
        transforms.ToTensor(),
        # transforms.Normalize(mean=[0.485, 0.456, 0.406],
        #                     std=[0.229, 0.224, 0.225])
    ])
    data = [pickle.loads(d) for d in data]
    # imgsName = glob.glob(path + "/*.jpg")
    # print(imgsName)
    imgs = [tf(i[1]).unsqueeze(0) for i in data]
    imgs = torch.cat(imgs, dim=0).to(torch.device(device))
    # print(imgs.shape)
    with torch.no_grad():
        y = model(imgs)

    return [[int(data[i][0]), y[i]] for i in range(len(data))]


def featureMapDistance(queryTensor: Tensor, cmpTensor: Tensor):

    distance = (queryTensor - cmpTensor).pow(2).sum(0)
    distance = float(distance)

    return distance


loadFeatureMapDict = {}
def loadFeatureMap(path: str, torchDevice:str="cuda"):
    if path in loadFeatureMapDict:
        return loadFeatureMapDict[path]

    while not os.path.exists(path):
        time.sleep(0.5)
    fm = torch.load(path, map_location=torch.device(torchDevice))
    loadFeatureMapDict[path] = fm

    return fm


def cleanLoadFeatureMapDict(l: list):
    for i in l:
        del loadFeatureMapDict[i[1]]


if __name__ == '__main__':
    model = torch.load("model_cpu.pt").to(torch.device("cuda"))
    res = getFeatureMapsToFMList(model, "cuda", r"D:\Projects\DeepLearning\reid\transport_car_6.0\CommonData\K28+020\20")
    print(res)